import torchvision
from torch.utils.data import DataLoader
from torchvision import transforms

from model import resnet50,resnet101

from torch import optim

from torch import nn

import torch
import os
import glob

transform = transforms.Compose([transforms.ToTensor(),transforms.Resize((224,224)),transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])])
train_set = torchvision.datasets.CIFAR10("../dataset",train=True,transform=transform,download=True)
test_set  = torchvision.datasets.CIFAR10("../dataset",train=False,transform=transform,download=True)

train_loader = DataLoader(train_set,batch_size=96)
test_loader = DataLoader(test_set,batch_size=96)


gpus = [0, 1]  # GPU索引
device = torch.device("cuda:0")



#模型定义
net = resnet101()
# net = net.to(device)
net = nn.DataParallel(net, device_ids=gpus)
net = net.to(gpus[0])  # 主模型


files = glob.glob(os.path.join("part_models_0728",f"*{'.pth'}"))
files = sorted(files, key=lambda x: int(x.split('.')[0].split("_")[-1]))
if len(files)>0:
    net=torch.load(files[-1])

#损失函数定义
lossFunc = nn.CrossEntropyLoss()
lossFunc = lossFunc.to(device)

#优化器
optimizer = optim.SGD(net.parameters(),lr=0.01)

os.makedirs("all_models101",exist_ok=True)
os.makedirs("part_models101",exist_ok=True)
f=open("accv3.0.txt","a+")
max_acc=0.65
accuracy = 0
start_epoch =0
for epoch in range(start_epoch,500):
    net.train()
    print("-------------第{}轮 训练-------------".format(epoch+1))
    for i,data in enumerate(train_loader):
        imgs,targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)
        output = net(imgs)
        # targets_pre = torch.argmax(output,dim=1)

        # print("imgs.shape:{},output.shape:{},targets.shape:{}".farmat(imgs.shape,output.shape,targets.shape))
        #
        # print("output:",output)
        # print("___________________")
        # print("targets:",targets)

        loss = lossFunc(output,targets)

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        if i%100==0:
            print("训练次数：{}，loss:{}".format(i+1,loss))

    net.eval()
    total_correct=0
    for data in test_loader:
        imgs,targets = data

        imgs = imgs.to(device)
        targets = targets.to(device)
        output = net(imgs)

        with torch.no_grad():
            labels_pre = torch.argmax(output,dim=1)
            total_correct += (labels_pre==targets).sum()
    accuracy = total_correct/len(test_set)
    print("accuracy:{}".format(accuracy))
    acc=round(accuracy.item(),2)
    f.write("acc[{}]:{}\n".format(epoch,acc))
    torch.save(net, os.path.join("all_models101","model_{}.pth".format(epoch)))
    if accuracy>max_acc:
        max_acc=accuracy
        torch.save(net, os.path.join("part_models101","model_{}.pth".format(epoch)))
        print("模型保存成功！")



















